import tensorflow as tf
from pathlib import Path
from abc import ABCMeta,abstractmethod
from stock_a.common.utils import os_util
from stock_a.feature import feature_resigteror

def sample_parse_func(example):
    label_keys = ['max_increase_20d', 'max_decrease_20d', 'max_increase_50d', 'max_decrease_50d', 'max_increase_120d', 'max_decrease_120d']
    feature_generator_class_list = feature_resigteror.get_feature_generator_class_list()
    feature_generator_list = [feature_generator_class(None, None) for feature_generator_class in feature_generator_class_list]
    feature_map = {}
    for feature_generator in feature_generator_list:
        feature_map[feature_generator.feature_name()] = feature_generator.generate_tf_type()
    for label_key in label_keys:
        feature_map[label_key] = tf.io.FixedLenFeature((), tf.float32)
    print(feature_map)
    parsed_example = tf.io.parse_single_example(example, features=feature_map)
    features = {k: v for k, v in parsed_example.items() if k not in label_keys}
    labels = {k: v for k, v in parsed_example.items() if k in label_keys}
    return features, labels

class BaseModel(metaclass=ABCMeta):
    def __init__(self, tf_model_class):
        self.quant_config = os_util.get_quant_config()
        tf_record_dir_path = self.quant_config['sample']['tf_record_dir_path']
        project_root = os_util.get_project_root()
        # build dataset
        tf_record_dir_path = project_root + tf_record_dir_path
        folder_path = Path(tf_record_dir_path)
        record_files = list(folder_path.glob('*.tfrecords'))
        batch_size = int(self.quant_config['train']['batch_size'])
        epochs = int(self.quant_config['train']['epochs'])
        pre_fetch_buffer = int(self.quant_config['train']['pre_fetch_buffer'])
        self.dataset = tf.data.TFRecordDataset([str(s) for s in record_files])
        self.dataset = self.dataset.map(map_func=sample_parse_func).prefetch(pre_fetch_buffer).batch(batch_size).repeat(epochs)
        # build model
        self.tf_model = tf_model_class()
        # parse ckpt dir
        self.ckpt_path = project_root + self.quant_config['train']['ckpt_path']
        # tensorboard
        self.tf_summary_writer = tf.summary.create_file_writer(project_root + self.quant_config['train']['tensorboard_log_dir'])

    @abstractmethod
    def train_model(self):
        raise Exception("BaseModel.train_model NotImplementedException")

    def save_checkpoint(self):
        self.tf_model.save_weights(self.ckpt_path)

    def load_checkpoint(self):
        self.tf_model.load_weights(self.ckpt_path)

    @abstractmethod
    def predict(self, inputs):
        raise Exception("BaseModel.predict NotImplementedException")
